#include "../Base.h"

EXTERN_C PSHORT NtBuildNumber;
static ULONG ObjectCallbackListOffset = 0;

LIST_ENTRY  g_OperationListHead;
FAST_MUTEX  g_OperationListLock;//�� g_OperationListHead  �ӵ�ͬ���ֶ� ������
PVOID       g_UpperHandle = NULL;
PVOID       g_LowerHandle = NULL;

#define DRIVER_TAG 'NS'

	
	
	BOOLEAN GetVersionAndHardCode()
	{
		BOOLEAN b = FALSE;
		switch (*NtBuildNumber)
		{
		case 7600:
		case 7601:
		{
			ObjectCallbackListOffset = 0xC0;
			b = TRUE;
			break;
		}
		case 9200:
		{
			ObjectCallbackListOffset = 0xC8;	//OBJECT_TYPE.CallbackList
			b = TRUE;
			break;
		}
		case 9600:
		{
			ObjectCallbackListOffset = 0xC8;	//OBJECT_TYPE.CallbackList
			b = TRUE;
			break;
		}
		default:  //win10  ����C8
			if (*NtBuildNumber > 10000)
			{
				ObjectCallbackListOffset = 0xc8;
				b = TRUE;
			}
			break;
		}
		return b;
	}


	//��ȡδ�ĵ����ĺ���
	PVOID GetUndocumentFunctionAddress(IN PUNICODE_STRING pFunName,
		IN PUCHAR pStartAddress,
		IN UCHAR* pFeatureCode,
		IN ULONG FeatureCodeNum,
		ULONG SerSize,
		UCHAR SegCode,
		ULONG AddNum,
		BOOLEAN ByName)
	{
		ULONG dwIndex = 0;
		PUCHAR pFunAddress = NULL;
		ULONG dwCodeNum = 0;

		if (pFeatureCode == NULL)
			return NULL;

		if (FeatureCodeNum >= 15)
			return NULL;

		if (SerSize > 0x1024)
			return NULL;

		if (ByName)
		{
			if (pFunName == NULL || !MmIsAddressValid(pFunName->Buffer))
				return NULL;

			pFunAddress = (PUCHAR)MmGetSystemRoutineAddress(pFunName);
			if (pFunAddress == NULL)
				return NULL;
		}
		else
		{
			if (pStartAddress == NULL || !MmIsAddressValid(pStartAddress))
				return NULL;

			pFunAddress = pStartAddress;
		}

		for (dwIndex = 0; dwIndex < SerSize; dwIndex++)
		{
			__try
			{
				if (pFunAddress[dwIndex] == pFeatureCode[dwCodeNum] || pFeatureCode[dwCodeNum] == SegCode)
				{
					dwCodeNum++;

					if (dwCodeNum == FeatureCodeNum)
						return pFunAddress + dwIndex - dwCodeNum + 1 + AddNum;

					continue;
				}

				dwCodeNum = 0;
			}
			__except (EXCEPTION_EXECUTE_HANDLER)
			{
				return 0;
			}
		}

		return 0;
	}

	//���㺯����ַ
	PVOID GetMovPoint(PVOID pCallPoint)
	{
		ULONG dwOffset = 0;
		ULONG_PTR returnAddress = 0;
		LARGE_INTEGER returnAddressTemp = { 0 };
		PUCHAR pFunAddress = NULL;

		if (pCallPoint == NULL || !MmIsAddressValid(pCallPoint))
			return NULL;

		pFunAddress = (PUCHAR)pCallPoint;
		// ����ƫ��  
		RtlCopyMemory(&dwOffset, (PVOID)(pFunAddress + 3), sizeof(ULONG));

		// JMP������ת  
		if ((dwOffset & 0x10000000) == 0x10000000)
		{
			dwOffset = dwOffset + 7 + (ULONG)pFunAddress;
			returnAddressTemp.QuadPart = (ULONG_PTR)pFunAddress & 0xFFFFFFFF00000000;
			returnAddressTemp.LowPart = dwOffset;
			returnAddress = returnAddressTemp.QuadPart;
			return (PVOID)returnAddress;
		}

		returnAddress = (ULONG_PTR)dwOffset + 7 + (ULONG)pFunAddress;
		return (PVOID)returnAddress;
	}


	PVOID GetPsLoadedListModule()
	{
		/*
		Ϊʲô��ʹ��DriverObjectȥö���� ��Ҫ��win10�����⵼���� ׷�沽���� ����
		RtlPcToFileHeader
		text:000000014015D6A3 48 8B 0D 86 50 0E 00                          mov     rcx, qword ptr cs:PsLoadedModuleList
		.text:000000014015D6AA 48 85 C9                                      test    rcx, rcx
		.text:000000014015D6AD 74 28                                         jz      short loc_14015D6D7
		.text:000000014015D6AF 48 8D 15 7A 50 0E 00                          lea     rdx, PsLoadedModuleList
		*/
		UNICODE_STRING usRtlPcToFileHeader = RTL_CONSTANT_STRING(L"RtlPcToFileHeader");
		UNICODE_STRING usPsLoadedModuleList = RTL_CONSTANT_STRING(L"PsLoadedModuleList");
		PVOID Point = NULL;
		static PVOID PsLoadedListModule = NULL;
		UCHAR shellcode[11] =
			"\x48\x8b\x0d\x60\x60\x60\x60"
			"\x48\x85\xc9";

		if (PsLoadedListModule)
			return PsLoadedListModule;

		if (*NtBuildNumber > 9600)
		{
			// win10 PsLoadedModuleList����
			PsLoadedListModule = MmGetSystemRoutineAddress(&usPsLoadedModuleList);
			return PsLoadedListModule;
		}

		Point = GetUndocumentFunctionAddress(&usRtlPcToFileHeader, NULL, shellcode, 10, 0xff, 0x60, 0, TRUE);
		if (Point == NULL || !MmIsAddressValid(Point))
			return NULL;

		Point = GetMovPoint(Point);
		if (Point == NULL || !MmIsAddressValid(Point))
			return NULL;

		PsLoadedListModule = Point;
		return PsLoadedListModule;
	}




	BOOLEAN ObGetDriverNameByPoint(ULONG_PTR Point/*��ַ*/, WCHAR* szDriverName)
	{
		PLDR_DATA_TABLE_ENTRY Begin = NULL;
		PLIST_ENTRY Head = NULL;  //����ڵ�
		PLIST_ENTRY Next = NULL;

		Begin = (PLDR_DATA_TABLE_ENTRY)GetPsLoadedListModule(); //��ȡģ������

		if (Begin == NULL)
			return FALSE;

		Head = (PLIST_ENTRY)Begin->InLoadOrderLinks.Flink;
		Next = Head->Flink;

		do
		{
			PLDR_DATA_TABLE_ENTRY Entry = CONTAINING_RECORD(Next, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
			Next = Next->Flink;

			if ((ULONG_PTR)Entry->DllBase <= Point && Point <= ((ULONG_PTR)Entry->DllBase + Entry->SizeOfImage))
			{
				if (szDriverName == NULL)
					return FALSE;
				RtlZeroMemory(szDriverName, 600);
				RtlCopyMemory(szDriverName, Entry->BaseDllName.Buffer, Entry->BaseDllName.Length);
				return TRUE;
			}

		} while (Next != Head->Flink);

		return FALSE;
	}


	ULONG EnumObRegisterCallBacks()
	{
		ULONG c = 0;
		PLIST_ENTRY CurrEntry = NULL;
		POB_CALLBACK pObCallback;
		ULONG64 ObProcessCallbackListHead = 0;
		ULONG64 ObThreadCallbackListHead = 0;
		WCHAR* szDriverBaseName = NULL;
		szDriverBaseName = (WCHAR*)ExAllocatePool(NonPagedPool, 600);
		if (szDriverBaseName == NULL)
			return FALSE;

		RtlZeroMemory(szDriverBaseName, 600);


		GetVersionAndHardCode();

		ObProcessCallbackListHead = ((ULONG64)(*(ULONG64*)PsProcessType) + ObjectCallbackListOffset); //����ģ��ص����� 
		ObThreadCallbackListHead = ((ULONG64)(*(ULONG64*)PsThreadType) + ObjectCallbackListOffset);;   //�߳�ģ��ص�����
		LOG("=====================================================���н��̻ص�====================================\r\n");
		CurrEntry = ((PLIST_ENTRY)ObProcessCallbackListHead)->Flink; //��õ�ǰ�ڵ�

		if (CurrEntry == NULL || !MmIsAddressValid(CurrEntry))
		{
			ExFreePool(szDriverBaseName);
			return 0;
		}

		do
		{
			pObCallback = (POB_CALLBACK)CurrEntry;
			if (pObCallback->ObHandle != 0)
			{
				if (ObGetDriverNameByPoint(pObCallback->PreCall, szDriverBaseName))
					LOG("DriverName=%S ObHandle=%p  attitude=%wZ PreCall=%p PostCall=%p \r\n",
						szDriverBaseName,
						pObCallback->ObHandle,
						&pObCallback->ObHandle->AltitudeString,
						pObCallback->PreCall,
						pObCallback->PostCall);




				c++;
			}

			CurrEntry = CurrEntry->Flink;

		} while (CurrEntry != (PLIST_ENTRY)ObProcessCallbackListHead);

		LOG("=====================================================�����̻߳ص�====================================\r\n");
		CurrEntry = ((PLIST_ENTRY)ObThreadCallbackListHead)->Flink;
		if (CurrEntry == NULL || !MmIsAddressValid(CurrEntry))
		{
			ExFreePool(szDriverBaseName);
			return c;
		}

		do
		{
			pObCallback = (POB_CALLBACK)CurrEntry;
			if (pObCallback->ObHandle != 0)
			{
				if (ObGetDriverNameByPoint(pObCallback->PreCall, szDriverBaseName))
					LOG("DriverName=%S ObHandle=%p  attitude=%wZ PreCall=%p PostCall=%p \r\n",
						szDriverBaseName,
						pObCallback->ObHandle,
						&pObCallback->ObHandle->AltitudeString,
						pObCallback->PreCall,
						pObCallback->PostCall);
				c++;
			}

			CurrEntry = CurrEntry->Flink;

		} while (CurrEntry != (PLIST_ENTRY)ObThreadCallbackListHead);

		LOG("=====================================================[�ص�����]:%x====================================\r\n",c);

		ExFreePool(szDriverBaseName);
		return c;
	}



	HANDLE GetCurrentProcessID()
	{
		return PsGetCurrentProcessId();
	}


	BOOLEAN IsMyProcess()
	{
		PEPROCESS Process = NULL;
		Process = PsGetCurrentProcess();
		//KeQuerySystemTime(1);
		if (_strnicmp("150330575.exe", (char*)PsGetProcessImageFileName(Process), strlen("150330575.exe")) == 0)
		{
			return TRUE;// OB_PREOP_SUCCESS;
		}
		else
		{
			return FALSE;
		}
	}



	//�ں˻�ȡϵͳ�������
	ULONG MyGetTickCount()
	{
		LARGE_INTEGER la;
		ULONG MyInc;
		MyInc = KeQueryTimeIncrement(); 


		KeQueryTickCount(&la);

		la.QuadPart *= MyInc;
		la.QuadPart /= 10000;

		return la.LowPart;

	}



	//APC_LEVEL=1;
	//�޸�Ȩ�ޣ������Ȩ��
	OB_PREOP_CALLBACK_STATUS Last_CallBack(IN PVOID RegistrationContext, IN POB_PRE_OPERATION_INFORMATION OperationInformation)
	{
		POPERATION_INFO_ENTRY NewEntry = NULL;


		if (PsGetCurrentProcess() == PsInitialSystemProcess)
			return OB_PREOP_SUCCESS;

		if (OperationInformation->ObjectType == *PsThreadType)
			return OB_PREOP_SUCCESS;


		if (IsMyProcess())
		{
			NewEntry = (POPERATION_INFO_ENTRY)ExAllocatePoolWithTag(NonPagedPool, sizeof(OPERATION_INFO_ENTRY), DRIVER_TAG);	//����Ȩ�ޱ���ռ�

			if (NewEntry)
			{
				NewEntry->Operation = OperationInformation->Operation;
				NewEntry->Flags = OperationInformation->Flags;
				NewEntry->Object = OperationInformation->Object;
				NewEntry->ObjectType = OperationInformation->ObjectType;
				//���������Ȩ�� ��ֹ�������ص� �޸� DuplicateHandle OpenProcess
				NewEntry->AccessMask = 0x1fffff;//OperationInformation->Parameters->CreateHandleInformation.DesiredAccess; /// Same for duplicate handle
				NewEntry->time = MyGetTickCount();// KeQuerySystemTime(1);

				ExAcquireFastMutex(&g_OperationListLock);//ͬ�����ٻ����� ��ȡ����Ȩ		
				InsertTailList(&g_OperationListHead, &NewEntry->ListEntry);//�����޸Ĺ��Ľڵ�
				LOG("============================================[��ATTITUDE]=====================================");
				LOG("[��ATTTITUDE]  [PID]:%llX  [��Ȩ��]:%llX [��Ȩ��]:%llx  [time]:%llX [line]:%lld\n",
					(ULONG64)PsGetCurrentProcessId(),
					(ULONG64)OperationInformation->Parameters->CreateHandleInformation.DesiredAccess,//ԭȨ�� ����
					(ULONG64)NewEntry->AccessMask,
					(ULONG64)NewEntry->time,
					(ULONG64)__LINE__);
				ExReleaseFastMutex(&g_OperationListLock);//ͬ�����ٻ����� �ͷ�����Ȩ
			}
		}

		UNREFERENCED_PARAMETER(RegistrationContext);

		return OB_PREOP_SUCCESS;
	}

	OB_PREOP_CALLBACK_STATUS First_CallBack(IN PVOID RegistrationContext, IN POB_PRE_OPERATION_INFORMATION OperationInformation)
	{
		PLIST_ENTRY ListEntry;
		UNREFERENCED_PARAMETER(RegistrationContext);

		if (IsMyProcess())
		{

			ExAcquireFastMutex(&g_OperationListLock);//��ȡͬ������Ȩ
			for (ListEntry = g_OperationListHead.Flink; ListEntry != &g_OperationListHead; ListEntry = ListEntry->Flink)
			{
				POPERATION_INFO_ENTRY Entry = (POPERATION_INFO_ENTRY)ListEntry;
				if (Entry->Operation == OperationInformation->Operation &&
					Entry->Flags == OperationInformation->Flags &&
					Entry->Object == OperationInformation->Object &&
					Entry->ObjectType == OperationInformation->ObjectType)
				{
					// 
					ULONG32 newAccess = OperationInformation->Parameters->CreateHandleInformation.DesiredAccess;
					ULONG32 oldAccess = Entry->AccessMask;
					LOG("============================================[��ATTITUDE]=====================================\r\n");
					LOG(" [��ATTTITUDE]  PID=%llX  [ԭȨ��]:%llX   [��Ȩ��]:%llX  [time]:llX  [line]:%d\n",
						(ULONG64)PsGetCurrentProcessId(),
						(ULONG64)(Entry->AccessMask),//ԭȨ��
						(ULONG64)newAccess,  //��Ȩ��
						(ULONG64)Entry->time,
						(ULONG64)__LINE__); 

					OperationInformation->Parameters->CreateHandleInformation.DesiredAccess = Entry->AccessMask;    //�ָ�Ȩ������OpenProcess
					OperationInformation->Parameters->DuplicateHandleInformation.DesiredAccess = Entry->AccessMask; //�ָ�Ȩ������ DuplicateHandle


					//�Ƴ����
					RemoveEntryList(&Entry->ListEntry);
					ExFreePoolWithTag(Entry, DRIVER_TAG);
					goto Release;

				}
			}
		Release:
			ExReleaseFastMutex(&g_OperationListLock);//�ͷ�ͬ������Ȩ
			//ok
		}

		return OB_PREOP_SUCCESS;
	}






	//ע��OB���ӳ�ʼ��
	OB_OPERATION_REGISTRATION ObUpperOperationRegistration[] =
	{
		{ NULL, OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE, Last_CallBack, NULL },//���̵�
		{ NULL, OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE, Last_CallBack, NULL },//�̵߳�
	};

	OB_OPERATION_REGISTRATION ObLowerOperationRegistration[] =
	{
		{ NULL, OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE, First_CallBack, NULL },//���̵�
		{ NULL, OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE, First_CallBack, NULL },//�̵߳�
	};

	// XignCode3 �ص��� 380800
	// EAC       �ص��� 327530
	// BE        �ص��� 363220
	OB_CALLBACK_REGISTRATION UpperCallbackRegistration =
	{
		OB_FLT_REGISTRATION_VERSION,
		2,
		RTL_CONSTANT_STRING(L"880000"),//�ߵ�
		NULL,
		ObUpperOperationRegistration
	};

	OB_CALLBACK_REGISTRATION LowerCallcackRegistration =
	{
		OB_FLT_REGISTRATION_VERSION,
		2,
		RTL_CONSTANT_STRING(L"10000"),//�͵�
		NULL,
		ObLowerOperationRegistration
	};

	void ObRegisterUnload()
	{
		if (NULL != g_LowerHandle)
			ObUnRegisterCallbacks(g_LowerHandle);
		if (NULL != g_UpperHandle)
			ObUnRegisterCallbacks(g_UpperHandle);
		while (!IsListEmpty(&g_OperationListHead))
			ExFreePoolWithTag(RemoveHeadList(&g_OperationListHead), DRIVER_TAG);
	}

	BOOLEAN ObRegisterCallBacksInit(PDRIVER_OBJECT pDriverObject)
	{

		LOG("============================================[��ʼ����OBJECT�ص�]======================================\r\n");
		NTSTATUS Status = STATUS_SUCCESS;
		PLDR_DATA ldr;

		ldr = (PLDR_DATA)pDriverObject->DriverSection;
		ldr->Flags |= 0x20;

		InitializeListHead(&g_OperationListHead);//��ʼ������ͷ
		ExInitializeFastMutex(&g_OperationListLock);//��ʼ�����ٻ�����

		ObUpperOperationRegistration[0].ObjectType = PsProcessType;
		//����ӵ�
		ObUpperOperationRegistration[0].Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;

		ObUpperOperationRegistration[1].ObjectType = PsThreadType;
		//����ӵ�
		ObUpperOperationRegistration[1].Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
		Status = ObRegisterCallbacks(&UpperCallbackRegistration, &g_UpperHandle);
		if (!NT_SUCCESS(Status))
		{
			g_UpperHandle = NULL;
			goto Exit;
		}

		ObLowerOperationRegistration[0].ObjectType = PsProcessType;
		//����ӵ�
		ObLowerOperationRegistration[0].Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;

		ObLowerOperationRegistration[1].ObjectType = PsThreadType;
		//����ӵ�
		ObLowerOperationRegistration[1].Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
		Status = ObRegisterCallbacks(&LowerCallcackRegistration, &g_LowerHandle);
		if (!NT_SUCCESS(Status))
		{
			g_LowerHandle = NULL;
			goto Exit;
		}

	Exit:
		if (!NT_SUCCESS(Status))
			ObRegisterUnload();

		return NT_SUCCESS(Status) ? TRUE : FALSE;
	}


